#include "OutputFileWriter.hh"
#include "SDBoxContainer.hh"
#include "YamlConfig.hh"
#include "TimeControl.hh"

#include <set>
#include <string>

OutputFileWriter::OutputFileWriter()
{
    m_TopMetalHist.clear();
    m_TopMetalTree.clear();
    m_iEvent = -1;

    YAML::Node outNode = YamlConfig::GetInstance().GetYaml()["Output"];
    m_TotalTruth = outNode["TotalTruth"].as<bool>(false);
    m_FileName = outNode["FileName"].as<std::string>("Digi.root");
    m_File = new TFile(m_FileName, "recreate");
    //m_SimEvent = new SimEvent();
    m_TruthTree = new TTree("TruthTree", "");
    m_TruthTree->Branch("TruthEnergy",              &m_TruthEnergy);
    m_TruthTree->Branch("TruthParticle",            &m_TruthParticle);
    m_TruthTree->Branch("TruthPosition",            &m_TruthPosition);
    m_TruthTree->Branch("TruthPolarization",        &m_TruthPolarization);
    m_TruthTree->Branch("TruthMomentumDirection",   &m_TruthMomentumDirection);
    m_TruthTree->Branch("TimeFlag",                 &m_TimeFlag);
    m_TruthTree->Branch("FirstSDStepTime",          &m_FirstSDStepTime);
    m_TruthTree->Branch("NumOfFiredSD",             &m_NumOfFiredSD,        "NumOfFiredSD/I");
    m_TruthTree->Branch("FiredSD",                  &m_FiredSD,             "FiredSD[NumOfFiredSD]/I");
    m_TruthTree->Branch("PolPhotElectronPos",       &m_PolPhotElectronPos);
    m_TruthTree->Branch("PolPhotElectronMom",       &m_PolPhotElectronMom);
    m_TruthTree->Branch("PolPhotElectronEkin",      &m_PolPhotElectronEkin);
    m_TruthTree->Branch("IfPolPhot",                &m_IfPolPhot);
    m_TruthTree->Branch("iEvent",                   &m_iEvent);
    m_TruthTree->Branch("IfPolPhotOnMCP",           &m_IfPolPhotOnMCP);
    m_TruthTree->Branch("NumOfDeltaE",              &m_NumOfDeltaE,         "NumOfDeltaE/I");
    m_TruthTree->Branch("DeltaEEnergy",             &m_DeltaEEnergy,        "DeltaEEnergy[NumOfDeltaE]/D");
    m_TruthTree->Branch("DeltaEPos",                &m_DeltaEPos,           "DeltaEPos[NumOfDeltaE][3]/D");
    m_TruthTree->Branch("DeltaEMom",                &m_DeltaEMom,           "DeltaEMom[NumOfDeltaE][3]/D");
    //if(m_TotalTruth == true)
    //{
    //    TBranch* br_SimEvent = m_TruthTree->Branch("SimEvent", "SimEvent", m_SimEvent, 32000, 0);
    //    br_SimEvent->SetAutoDelete(true);
    //}
}

OutputFileWriter::~OutputFileWriter()
{
    for (auto MapHist : m_TopMetalHist)
    {
        delete MapHist;
        MapHist = nullptr;
    }
    m_File->Close();
}

TH2F* OutputFileWriter::CreateTopMetalTree(UInt_t iTopMetal)
{
    if(iTopMetal != m_TopMetalTree.size())
    {
        std::cerr << "Input iTopMetal and Number of exist TopMetalTree not matched!!!" << std::endl;
        exit(EXIT_FAILURE);
    }
    if(iTopMetal >= MAXSD)
    {
        std::cerr << "OutputFileWriter: Number of SD overflow!!! Please enlarge MAXSD in OutputFileWriter.hh" << std::endl;

    }
    std::pair<UInt_t, UInt_t> numOfPixel = SDBoxContainer::GetInstance()[iTopMetal].NumOfPixel();
    m_File->cd();
    m_TopMetalTree.push_back(new TTree(TString::Format("TopMetal_%d", iTopMetal), ""));
    TH2F* th_TopMetal = new TH2F(TString::Format("th2_TopMetal_%d", iTopMetal), "", numOfPixel.first, 0, numOfPixel.first+1, numOfPixel.second, 0, numOfPixel.second+1);
    m_TopMetalHist.push_back(th_TopMetal);
    m_TopMetalTree[iTopMetal]->Branch("FrameWriteTime", &(m_FrameWriteTime[iTopMetal]));
    m_TopMetalTree[iTopMetal]->Branch("TopMetalHist", m_TopMetalHist[iTopMetal]);
    m_TopMetalTree[iTopMetal]->Branch("iEvent", &m_iEvent);
    m_TopMetalTree[iTopMetal]->Branch("iFrame", &m_iFrame[iTopMetal]);
    m_TopMetalTree[iTopMetal]->Branch("iGMCP", &m_iGMCP[iTopMetal]);

    m_GMCPTree.push_back(new TTree(TString::Format("GMCP_%d", iTopMetal), ""));
    m_GMCPTree[iTopMetal]->Branch("iEvent", &m_iEvent);
    m_GMCPTree[iTopMetal]->Branch("iFrame", &m_iFrame[iTopMetal]);
    m_GMCPTree[iTopMetal]->Branch("GMCPTime", &m_GMCPTime[iTopMetal]);
    m_GMCPTree[iTopMetal]->Branch("GMCPAmp", &m_GMCPAmp[iTopMetal]);
    m_iGMCP[iTopMetal] = 0;

    return m_TopMetalHist[iTopMetal];
}

void OutputFileWriter::WriteTruthInfo(SimEvent* aEvent, TruthInfo* truthInfo)
{
    //std::cout << "Write Truth Info at Time " << FirstSDStepTime << std::endl;
    m_TruthEnergy = aEvent->GetTruthEnergy();
    m_TruthParticle = aEvent->GetTruthParticle();
    m_TruthPosition = aEvent->GetTruthPosition();
    m_TruthPolarization = aEvent->GetTruthPolarization();
    m_TruthMomentumDirection = aEvent->GetTruthMomentumDirection();
    m_TimeFlag = aEvent->GetTimeFlag();
    m_FirstSDStepTime = truthInfo->FirstSDStepTime;

    //if(m_TotalTruth == true)
    //{
    //    m_SimEvent->MyClear();
    //    //(*m_SimEvent) = *aEvent;
    //    m_SimEvent->CopyEvent(aEvent);
    //std::cout << __LINE__ << std::endl;
    //}

    m_NumOfFiredSD = 0;
    for (auto iSD : truthInfo->FiredSD)
    {
        m_FiredSD[m_NumOfFiredSD] = iSD;
        ++m_NumOfFiredSD;
    }
    m_IfPolPhot = truthInfo->IfPolPhot;
    m_IfPolPhotOnMCP = truthInfo->IfPolPhotOnMCP;
    m_PolPhotElectronPos = truthInfo->PolPhotElectronPos;
    m_PolPhotElectronMom = truthInfo->PolPhotElectronMom;
    m_PolPhotElectronEkin = truthInfo->PolPhotElectronEkin;

    m_NumOfDeltaE = truthInfo->DeltaEList.size();
    if (m_NumOfDeltaE > MAXDeltaE)
    {
        std::cerr << "Error! Number of Delta Electrons larger than MAXDeltaE. Please enlarge it." << std::endl;
        exit(-1);
    }
    for (int ie = 0; ie < m_NumOfDeltaE; ++ie)
    {
        m_DeltaEEnergy[ie] = truthInfo->DeltaEList[ie].Energy;
        for (int iarr = 0; iarr < 3; ++iarr)
        {
            m_DeltaEPos[ie][iarr] = truthInfo->DeltaEList[ie].Position[iarr];
            m_DeltaEMom[ie][iarr] = truthInfo->DeltaEList[ie].Mom[iarr];
        }
    }

    ++m_iEvent;
    m_TruthTree->Fill();
    //m_SimEvent->MyClear();
}

void OutputFileWriter::WriteGMCP(UInt_t iTopMetal, Double_t GMCPTime, Double_t GMCPAmp)
{
    m_GMCPTime[iTopMetal] = GMCPTime;
    m_GMCPAmp[iTopMetal] = GMCPAmp;

    m_GMCPTree[iTopMetal]->Write();
    ++m_iGMCP[iTopMetal];
}

void OutputFileWriter::WriteFrame(UInt_t iTopMetal, Double_t FrameWriteTime, UInt_t iFrame)
{
    //std::cout << "Write Frame of TopMetal " << iTopMetal << std::endl;
    m_FrameWriteTime[iTopMetal] = FrameWriteTime;
    m_iFrame[iTopMetal] = iFrame;

    m_TopMetalTree[iTopMetal]->Fill();
    if(TimeControl::Verbose())
    {
        std::cout << "OutputFileWriter::WriteFrame(): TopMetal: " << iTopMetal << "\tFrame: " << iFrame << "\tTime: " << FrameWriteTime << std::endl;
    }
}

void OutputFileWriter::SaveAll()
{
    if(TimeControl::Verbose()) std::cout << "OutputFileWriter::SaveAll(): Save all info." << std::endl;
    m_File->cd();
    //m_ConfigTree->AutoSave();
    m_TruthTree->AutoSave();
    for(auto iTMTree:m_TopMetalTree)
    {
        m_File->cd();
        iTMTree->AutoSave();
    }
    for(auto iGMCPTree:m_GMCPTree)
    {
        m_File->cd();
        iGMCPTree->AutoSave();
    }
}